#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/correlator.h>
#include <src/operator.h>
#include <src/cluster.h>
#include <src/expectation_value.h>
#include <src/blas_lapack.h>
#include <src/parallel.h>
#include <src/random.h>


template <class S>
void pcps::update_correlator(const pcps::Input & userinp,
                             const pcps::Hamiltonian<S> & ham,
                             const pcps::Wavefunction<S> & wfn,
                             pcps::Correlator<S> & corr,
                             const double noise_magnitude,
                             const double damping,
                             const int perturbation_level)
{

  // define the worker class that will compute the local hamiltonian and overlap matrices in parallel
  class Local_Worker : public pcps::Parallel_Worker_Base<S> {

    public:

      pcps::Correlator<S> _corr;

      void do_work(const pcps::Input & userinp,
                   const pcps::Hamiltonian<S> & ham,
                   const pcps::Wavefunction<S> & wfn,
                   const int perturbation_level,
                   const boost::shared_array<S> & output_array,
                   const int array_size,
                   const int my_min,
                   const int my_max,
                   const std::vector<int> & slaves,
                   int & slaves_used,
                   int & jobs_in_progress)
      {

        // get MPI info
        const MPI::Comm & comm = MPI::COMM_WORLD;
        const int nproc = comm.Get_size();
        const int myrank = comm.Get_rank();

        // prepare a copy of the identity operator
        const typename pcps::OpBase<S>::const_ptr identity_op(new pcps::OpIdentity<S>(pcps::unity<S>()));

        // exclude the correlator being updated from being used when calculating expectation values
        typename pcps::Correlator<S>::Set excluded;
        excluded.insert(_corr);

        // compute elements of the local Hamiltonian and overlap matrices
        for (pcps::Cluster bra_occs(userinp.nsites(), userinp.configs_per_site(), _corr.site_vec()); !bra_occs.finished(); bra_occs++)
        for (pcps::Cluster ket_occs(userinp.nsites(), userinp.configs_per_site(), _corr.site_vec()); !ket_occs.finished(); ket_occs++) {

          // get the local hamiltonian compound index
          const int cmpd_h_bra_index = bra_occs.cmpd();
          const int cmpd_h_ket_index = ket_occs.cmpd();
          const int cmpd_h_index = _corr.nelements() * cmpd_h_bra_index + cmpd_h_ket_index;

          // compute each operator's contribution to the local hamiltonian element
          for (int i = my_min; i < my_max; i++)
            pcps::expectation_value<S>(slaves, slaves_used, jobs_in_progress, output_array.get(), cmpd_h_index, pcps::unity<S>(),
                                            userinp, wfn, ham.ops().at(i), perturbation_level, bra_occs, ket_occs, excluded);

          // if the bra and ket configurations are the same, compute the overlap matrix element
          if (myrank == 0 && cmpd_h_bra_index == cmpd_h_ket_index) {
            const int array_pos = _corr.nelements() * _corr.nelements() + cmpd_h_bra_index;
            pcps::expectation_value<S>(slaves, slaves_used, jobs_in_progress, output_array.get(), array_pos, pcps::unity<S>(),
                                            userinp, wfn, identity_op, perturbation_level, bra_occs, ket_occs, excluded);
          }

        }

      }

  };

  // create the worker that will compute the local hamiltonian and overlap matrices
  Local_Worker worker;

  // give the correlator we are updating to the worker (shallow copy)
  worker._corr = corr;

  // compute the local hamiltonian and overlap matrices in parallel
  const int n = corr.nelements();
  boost::shared_array<S> data_array = pcps::compute_in_parallel(userinp, ham, wfn, perturbation_level, &worker, n*n+n);

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // diagonalize the local hamiltonian and compute the new correlator values on the root node
  if (myrank == 0) {

    // create pointers for the local Hamiltonian and overlap matrices
    S * const local_h = &data_array[0];
    S * const local_s = &data_array[n*n];

    // print the overlap matrix
    if (userinp.verbose_print()) {
      for (pcps::Cluster occs(userinp.nsites(), userinp.configs_per_site(), corr.site_vec()); !occs.finished(); occs++) {
        std::cout << boost::format("for occupation (");
        for (int p = 0; p < occs.size(); p++)
          std::cout << boost::format(" %i") % occs.ket_occ(occs.sites().at(p));
        std::cout << boost::format(" ) overlap eigenvalue = %s") % pcps::format_number("%10.4e", local_s[occs.cmpd()]) << std::endl;
      }
      std::cout << std::endl;
    }

    // print the local hamiltonian matrix
    if (userinp.verbose_print()) {
      std::cout << boost::format(" local hamiltonian matrix:") << std::endl;
      for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++)
          std::cout << pcps::format_number(" %10.4f", local_h[i*n + j]);
        std::cout << std::endl;
      }
      std::cout << std::endl;
    }

    // divide the local hamiltonian matrix by the overlap matrix
    for (int i = 0; i < n; i++) {
      if (std::abs(local_s[i]) < 1.0e-6)
        local_s[i] = pcps::unity<S>();
        //throw pcps::Exception("small overlap eigenvalue detected");
      for (int j = 0; j < n; j++)
        local_h[i*n + j] /= local_s[i];
    }

    // prepare variables for lapack routine
    const int lwork = 10*n;
    int info;
    boost::shared_array<S> A     = pcps::allocate_shared_array<S>(n*n);
    boost::shared_array<S> VR    = pcps::allocate_shared_array<S>(n*n);
    boost::shared_array<S> w     = pcps::allocate_shared_array<S>(n);
    boost::shared_array<double> wr    = pcps::allocate_shared_array<double>(n);
    boost::shared_array<double> wi    = pcps::allocate_shared_array<double>(n);
    boost::shared_array<S> work  = pcps::allocate_shared_array<S>(lwork);
    boost::shared_array<double> rwork = pcps::allocate_shared_array<double>(2*n);
    S * VL; // not used

    // transpose the local hamiltonian matrix so that we have fortran's column major ordering
    for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      A[j*n+i] = local_h[i*n+j];

    // find the eigenvalues and right eigenvectors of the local hamiltonian
    pcps::xgeev('N', 'V', n, A.get(), n, w.get(), wr.get(), wi.get(), VL, n, VR.get(), n, work.get(), lwork, rwork.get(), info);
    if (info != 0)
      throw pcps::Exception( (boost::format("xgeev failed for local hamiltonian with error code %i") % info).str() );

    // the p-th eigenvector is now stored in VR[p,*]

    // find the lowest eigenvalue
    int p = -1;
    {
      double min_eval = 1.0e10;
      for (int i = 0; i < n; i++) {
        if (userinp.verbose_print())
          std::cout << boost::format(" eval %2i = %11.4e + %11.4e * i") % i % wr[i] % wi[i] << std::endl;
        if (wr[i] < min_eval) {
          min_eval = wr[i];
          p = i;
        }
      }
      if (p > 0 && std::fabs(wr[p]-wr[p-1]) < 1.0e-10 && std::fabs(wi[p]+wi[p-1]) < 1.0e-10)
        p = p - 1;
    }
    if (p == -1)
      throw pcps::Exception( (boost::format("no acceptable eigenvalues")).str() );

    if (userinp.verbose_print()) {
      std::cout << boost::format("chose eval %2i = %11.4e") % p % wr[p] << std::endl;
      std::cout << std::endl;
    }

    // print the eigenvectors
    if (userinp.verbose_print()) {
      std::cout << boost::format(" right eigenvectors:") << std::endl;
      for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++)
          std::cout << pcps::format_number(" %10.4f", VR[i*n+j]);
        std::cout << std::endl;
      }
      std::cout << std::endl;
    }

    // update the correlator, setting its values equal to the eigenvector with lowest eigenvalue
    S max_value = 0.00;
    S min_value = 1.0e100;
    for (int i = 0; i < n; i++) {
      corr[i] = damping * corr[i] + ( 1.0 - damping ) * VR[p*n+i];
      if (std::abs(corr[i]) > std::abs(max_value))
        max_value = corr[i];
      if (std::abs(corr[i]) < std::abs(min_value))
        min_value = corr[i];
    }
    for (int i = 0; i < n; i++)
      corr[i] = corr[i] / max_value;
    //if ( ( max_value > 0 && min_value > 0 ) || ( max_value < 0 && min_value < 0 ) )
    //  for (int i = 0; i < n; i++)
    //    corr[i] = corr[i] / ( (max_value + min_value) / 2.0 );
    //else
    //  for (int i = 0; i < n; i++)
    //    corr[i] = corr[i] / max_value;

    // add noise
    for (int i = 0; i < corr.nelements(); i++)
      if (typeid(S) == typeid(double))
        corr[i] += noise_magnitude * 2.0 * ( pcps::random_number<S>() - 0.5 );
      else if (typeid(S) == typeid(std::complex<double>))
        corr[i] += noise_magnitude * 2.0 * (   pcps::random_number<S>()
                                             - 0.5 * pcps::unity<S>()
                                             - 0.5 * std::sqrt( -1.0 * pcps::unity<S>() )
                                           );
      else
        throw pcps::Exception("unknown S type in pcps::update_correlator");

    //// print the largest deviation from 1 for perturbative correlators
    //if (perturbative) {
    //  double max_deviation = 0.0;
    //  for (int i = 0; i < n; i++)
    //    if (std::fabs(corr[i]) > max_deviation)
    //      max_deviation = std::fabs(corr[i]);
    //  std::cout << boost::format("maximum deviation = %10.6f") % max_deviation << std::endl;
    //  std::cout << std::endl;
    //}

  }

  // broadcast the new correlator values to all processes
  comm.Barrier();
  comm.Bcast((void *)corr.data_ptr(), sizeof(S) * corr.nelements(), MPI::CHAR, 0);

//  // get MPI info
//  const MPI::Comm & comm = MPI::COMM_WORLD;
//  const int nproc = comm.Get_size();
//  const int myrank = comm.Get_rank();
//
//  // set some constants
//  const int n_lattice_sites = userinp.nsites();
//  const int configs_per_site = corr.configs_per_site();
//  const int n_corr_sites = corr.sites().size();
//
//  // determine the type of parallelism to use
//  const bool site_parallel = ( userinp.parallel_type() == pcps::keyword::Parallel_Site );
//  const bool farm_parallel = ( userinp.parallel_type() == pcps::keyword::Parallel_Farm );
//
//  // exclude the correlator being updated from being used in calculating expectation values
//  typename pcps::Correlator<S>::Set excluded;
//  excluded.insert(corr);
//
//  // allocate vectors to hold the local hamiltonian and overlap matrices
//  const int corr_configs = pcps::pow<int>(configs_per_site, n_corr_sites);
//  const int array_size = corr_configs * corr_configs + corr_configs;
//  boost::shared_array<S> data_array = pcps::allocate_shared_array<S>( (myrank == 0 ? array_size : 1), pcps::zero<S>() );
//  boost::shared_array<S> data_array_dist = (site_parallel ? pcps::allocate_shared_array<S>(array_size, pcps::zero<S>()) : data_array);
//
//  // prepare a copy of the identity operator
//  const pcps::OpBase_const_ptr identity_op(new pcps::OpIdentity(1.0));
//
//  // In farm parallel mode, only the master process deals with setting up expectation values.
//  // In site parallel mode, each processor computes the expectation values for a subset of the hamiltonian's operators.
//  if (myrank == 0 || site_parallel) {
//
//    // variables to keep track of what the slaves are doing in farm parallel mode
//    std::vector<int> slaves(nproc-1);
//    for (int i = 0; i < slaves.size(); i++)
//      slaves.at(i) = i+1;
//    int slaves_used = 0;
//    int jobs_in_progress = 0;
//
//    // determine which of the hamiltonian operators belong to this process
//    const int extras = ham.ops().size() % nproc;
//    const int my_min = ( site_parallel ? myrank * ( ham.ops().size() / nproc ) + std::min(myrank, extras) : 0 );
//    const int my_max = ( site_parallel ? ( myrank + 1 ) * ( ham.ops().size() / nproc ) + std::min(myrank + 1, extras) : ham.ops().size() );
//
//    // compute elements of the local Hamiltonian and overlap matrices
//    for (pcps::Cluster bra_occs(n_lattice_sites, configs_per_site, corr.sites()); !bra_occs.finished(); bra_occs++)
//    for (pcps::Cluster ket_occs(n_lattice_sites, configs_per_site, corr.sites()); !ket_occs.finished(); ket_occs++) {
//
//      // get the local hamiltonian compound index
//      const int cmpd_h_bra_index = bra_occs.cmpd();
//      const int cmpd_h_ket_index = ket_occs.cmpd();
//      const int cmpd_h_index = corr_configs * cmpd_h_bra_index + cmpd_h_ket_index;
//
//      // compute each operator's contribution to the local hamiltonian element
//      for (int i = my_min; i < my_max; i++)
//        pcps::expectation_value<S>(slaves, slaves_used, jobs_in_progress, data_array_dist.get(), cmpd_h_index, pcps::unity<S>(),
//                                        userinp, wfn, ham.ops().at(i), perturbation_level, bra_occs, ket_occs, excluded);
//
//      // if the bra and ket configurations are the same, compute the overlap matrix element
//      if (myrank == 0 && bra_occs == ket_occs) {
//        const int array_pos = corr_configs * corr_configs + cmpd_h_bra_index;
//        pcps::expectation_value<S>(slaves, slaves_used, jobs_in_progress, data_array_dist.get(), array_pos, pcps::unity<S>(),
//                                        userinp, wfn, identity_op, perturbation_level, bra_occs, ket_occs, excluded);
//      }
//
//    }
//
//    // if working in site parallel mode, collect the local hamiltonian elements on the root node
//    if (site_parallel) {
//      comm.Barrier();
//      if (typeid(S) == typeid(double))
//        comm.Reduce((void *)(&data_array_dist[0]), (void *)(&data_array[0]), array_size, MPI::DOUBLE, MPI::SUM, 0);
//      else if (typeid(S) == typeid(std::complex<double>))
//        comm.Reduce((void *)(&data_array_dist[0]), (void *)(&data_array[0]), array_size, MPI::DOUBLE_COMPLEX, MPI::SUM, 0);
//      else
//        throw pcps::Exception("unknown S type in pcps::update_correlator");
//    }
//
//    // if working in farm parallel mode, collect results from any slaves that are still working
//    else if (farm_parallel)
//      pcps::collect_expectation_values(slaves, jobs_in_progress, &data_array[0]);
//
//    else
//      throw pcps::Exception("unknown parallelism type");
//
//
//    // diagonalize the local hamiltonian and compute the new correlator values on the root node
//    if (myrank == 0) {
//
//      // create pointers for the local Hamiltonian and overlap matrices
//      S * const local_h = &data_array[0];
//      S * const local_s = &data_array[corr_configs * corr_configs];
//
//      // print the overlap matrix
//      if (userinp.verbose_print()) {
//        for (pcps::Cluster occs(n_lattice_sites, configs_per_site, corr.sites()); !occs.finished(); occs++) {
//          std::cout << boost::format("for occupation (");
//          for (int p = 0; p < occs.size(); p++)
//            std::cout << boost::format(" %i") % occs.ket_occ(occs.sites().at(p));
//          std::cout << boost::format(" ) overlap eigenvalue = %s") % pcps::format_number("%10.4e", local_s[occs.cmpd()]) << std::endl;
//        }
//        std::cout << std::endl;
//      }
//
//      // print the local hamiltonian matrix
//      if (userinp.verbose_print()) {
//        std::cout << boost::format(" local hamiltonian matrix:") << std::endl;
//        for (int i = 0; i < corr_configs; i++) {
//          for (int j = 0; j < corr_configs; j++)
//            std::cout << pcps::format_number(" %10.4f", local_h[i*corr_configs + j]);
//          std::cout << std::endl;
//        }
//        std::cout << std::endl;
//      }
//
//      // divide the local hamiltonian matrix by the overlap matrix
//      for (int i = 0; i < corr_configs; i++) {
//        if (std::abs(local_s[i]) < 1.0e-6)
//          local_s[i] = pcps::unity<S>();
//          //throw pcps::Exception("small overlap eigenvalue detected");
//        for (int j = 0; j < corr_configs; j++)
//          local_h[i*corr_configs + j] /= local_s[i];
//      }
//
//      // prepare variables for lapack routine
//      const int n = corr_configs;
//      const int lwork = 10*n;
//      int info;
//      boost::shared_array<S> A     = pcps::allocate_shared_array<S>(n*n);
//      boost::shared_array<S> VR    = pcps::allocate_shared_array<S>(n*n);
//      boost::shared_array<S> w     = pcps::allocate_shared_array<S>(n);
//      boost::shared_array<double> wr    = pcps::allocate_shared_array<double>(n);
//      boost::shared_array<double> wi    = pcps::allocate_shared_array<double>(n);
//      boost::shared_array<S> work  = pcps::allocate_shared_array<S>(lwork);
//      boost::shared_array<double> rwork = pcps::allocate_shared_array<double>(2*n);
//      S * VL; // not used
//
//      // transpose the local hamiltonian matrix so that we have fortran's column major ordering
//      for (int i = 0; i < n; i++)
//      for (int j = 0; j < n; j++)
//        A[j*n+i] = local_h[i*n+j];
//
//      // find the eigenvalues and right eigenvectors of the local hamiltonian
//      pcps::xgeev('N', 'V', n, A.get(), n, w.get(), wr.get(), wi.get(), VL, n, VR.get(), n, work.get(), lwork, rwork.get(), info);
//      if (info != 0)
//        throw pcps::Exception( (boost::format("xgeev failed for local hamiltonian with error code %i") % info).str() );
//
//      // the p-th eigenvector is now stored in VR[p,*]
//
//      // find the lowest eigenvalue
//      int p = -1;
//      {
//        double min_eval = 1.0e10;
//        for (int i = 0; i < n; i++) {
//          if (userinp.verbose_print())
//            std::cout << boost::format(" eval %2i = %11.4e + %11.4e * i") % i % wr[i] % wi[i] << std::endl;
//          if (wr[i] < min_eval) {
//            min_eval = wr[i];
//            p = i;
//          }
//        }
//        if (p > 0 && std::fabs(wr[p]-wr[p-1]) < 1.0e-10 && std::fabs(wi[p]+wi[p-1]) < 1.0e-10)
//          p = p - 1;
//      }
//      if (p == -1)
//        throw pcps::Exception( (boost::format("no acceptable eigenvalues")).str() );
//
//      if (userinp.verbose_print()) {
//        std::cout << boost::format("chose eval %2i = %11.4e") % p % wr[p] << std::endl;
//        std::cout << std::endl;
//      }
//
//      // print the eigenvectors
//      if (userinp.verbose_print()) {
//        std::cout << boost::format(" right eigenvectors:") << std::endl;
//        for (int i = 0; i < n; i++) {
//          for (int j = 0; j < n; j++)
//            std::cout << pcps::format_number(" %10.4f", VR[i*n+j]);
//          std::cout << std::endl;
//        }
//        std::cout << std::endl;
//      }
//
//      // update the correlator, setting its values equal to the eigenvector with lowest eigenvalue
//      S max_value = 0.00;
//      S min_value = 1.0e100;
//      for (int i = 0; i < n; i++) {
//        corr[i] = damping * corr[i] + ( 1.0 - damping ) * VR[p*n+i];
//        if (std::abs(corr[i]) > std::abs(max_value))
//          max_value = corr[i];
//        if (std::abs(corr[i]) < std::abs(min_value))
//          min_value = corr[i];
//      }
//      for (int i = 0; i < n; i++)
//        corr[i] = corr[i] / max_value;
//      //if ( ( max_value > 0 && min_value > 0 ) || ( max_value < 0 && min_value < 0 ) )
//      //  for (int i = 0; i < n; i++)
//      //    corr[i] = corr[i] / ( (max_value + min_value) / 2.0 );
//      //else
//      //  for (int i = 0; i < n; i++)
//      //    corr[i] = corr[i] / max_value;
//
//      // add noise
//      for (int i = 0; i < corr.nelements(); i++)
//        if (typeid(S) == typeid(double))
//          corr[i] += noise_magnitude * 2.0 * ( pcps::random_number<S>() - 0.5 );
//        else if (typeid(S) == typeid(std::complex<double>))
//          corr[i] += noise_magnitude * 2.0 * (   pcps::random_number<S>()
//                                               - 0.5 * pcps::unity<S>()
//                                               - 0.5 * std::sqrt( -1.0 * pcps::unity<S>() )
//                                             );
//        else
//          throw pcps::Exception("unknown S type in pcps::update_correlator");
//
//      //// print the largest deviation from 1 for perturbative correlators
//      //if (perturbative) {
//      //  double max_deviation = 0.0;
//      //  for (int i = 0; i < n; i++)
//      //    if (std::fabs(corr[i]) > max_deviation)
//      //      max_deviation = std::fabs(corr[i]);
//      //  std::cout << boost::format("maximum deviation = %10.6f") % max_deviation << std::endl;
//      //  std::cout << std::endl;
//      //}
//
//    }
//
//  // slave processes in farm parallel mode
//  } else if (farm_parallel) {
//
//    // wait for instructions from the master process
//    pcps::expectation_value_slave(userinp, wfn);
//
//  } else
//    throw pcps::Exception("unknown parallelism type");
//
//  // broadcast the new correlator values to all processes
//  comm.Barrier();
//  comm.Bcast((void *)corr.data_ptr(), sizeof(S) * corr.nelements(), MPI::CHAR, 0);

}

// instantiate the template
template void pcps::update_correlator(const pcps::Input & userinp,
                                      const pcps::Hamiltonian<double> & ham,
                                      const pcps::Wavefunction<double> & wfn,
                                      pcps::Correlator<double> & corr,
                                      const double noise_magnitude,
                                      const double damping,
                                      const int perturbation_level);
template void pcps::update_correlator(const pcps::Input & userinp,
                                      const pcps::Hamiltonian< std::complex<double> > & ham,
                                      const pcps::Wavefunction< std::complex<double> > & wfn,
                                      pcps::Correlator< std::complex<double> > & corr,
                                      const double noise_magnitude,
                                      const double damping,
                                      const int perturbation_level);
